package org.jetbrains.kotlinx.dataframe.plugin.impl.api

import org.jetbrains.kotlinx.dataframe.plugin.extensions.KotlinTypeFacade
import org.jetbrains.kotlinx.dataframe.plugin.extensions.changeNullability
import org.jetbrains.kotlinx.dataframe.plugin.extensions.isList
import org.jetbrains.kotlinx.dataframe.plugin.extensions.typeArgument
import org.jetbrains.kotlinx.dataframe.plugin.impl.*
import org.jetbrains.kotlinx.dataframe.plugin.impl.data.ColumnWithPathApproximation

internal class Explode0 : AbstractInterpreter<PluginDataFrameSchema>() {
    val Arguments.dropEmpty: Boolean by arg(defaultValue = Present(true))
    val Arguments.receiver: PluginDataFrameSchema by dataFrame()
    val Arguments.selector: ColumnsResolver? by arg(defaultValue = Present(null))
    override val Arguments.startingSchema get() = receiver

    override fun Arguments.interpret(): PluginDataFrameSchema {
        val columns = selector ?: object : ColumnsResolver {
            override fun resolve(df: PluginDataFrameSchema): List<ColumnWithPathApproximation> {
                return df.flatten(includeFrames = false).filter {
                    val column = it.column
                    column is SimpleFrameColumn || column is SimpleDataColumn && column.type.isList()
                }
            }
        }
        return receiver.explodeImpl(dropEmpty, columns.resolve(receiver))
    }
}

context(facade: KotlinTypeFacade)
fun PluginDataFrameSchema.explodeImpl(dropEmpty: Boolean, selector: List<ColumnWithPathApproximation>): PluginDataFrameSchema {
    val selected = selector.associateBy { it.path }

    fun makeNullable(column: SimpleCol): SimpleCol {
        return when (column) {
            is SimpleColumnGroup -> SimpleColumnGroup(column.name, column.columns().map { makeNullable(it) })
            is SimpleFrameColumn -> column
            is SimpleDataColumn -> {
                column.changeType(type = column.type.changeNullability { nullable -> selector.size > 1 || !dropEmpty || nullable })
            }
        }
    }

    fun explode(column: SimpleCol, path: List<String>): SimpleCol {
        val fullPath = path + listOf(column.name)
        return when (column) {
            is SimpleColumnGroup -> {
                SimpleColumnGroup(column.name, column.columns().map { explode(it, fullPath) })
            }
            is SimpleFrameColumn -> {
                val s = selected[fullPath]
                if (s != null) {
                    SimpleColumnGroup(s.column.name, column.columns().map { makeNullable(it) })
                } else {
                    column
                }
            }
            is SimpleDataColumn -> {
                val s = selected[fullPath]
                if (s != null) {
                    val newType = when {
                        column.type.isList() -> column.type.typeArgument()
                        else -> column.type
                    }
                    makeNullable(simpleColumnOf(s.column.name, newType.coneType))
                } else {
                    column
                }
            }
        }
    }

    return PluginDataFrameSchema(
        columns().map { column ->
            explode(column, emptyList())
        }
    )
}
